import base64
import uuid

from openpyxl import load_workbook
from PIL import Image
import io
from AiUtil import extract_img_feature
import sqlite3


def get_images_by_row(ws):
    images_dict = {}
    for img in ws._images:  # 获取所有图片对象
        anchor_cell = img.anchor._from  # 获取图片锚定的第一个单元格
        row = anchor_cell.row  # 获取行号
        if row not in images_dict:
            images_dict[row] = []
        images_dict[row].append(img)  # 将图片添加到对应行的列表中

    return images_dict


def load_excel_file():
    # 加载现有的 Excel 文件
    wb = load_workbook('data/笔记词典/附件下载_寒春笔记词典众人拾柴版/寒春笔记词典众人拾柴版.xlsx')
    ws = wb.active  # 或者指定工作表名称 wb['SheetName']

    images_by_row = get_images_by_row(ws)

    item_list = []
    for row_idx, row in enumerate(ws.iter_rows(values_only=True), start=0):
        print(f"Processing row {row_idx}")
        if row_idx not in images_by_row:
            continue

        img_data = images_by_row[row_idx][0]._data()
        pil_img = Image.open(io.BytesIO(img_data))
        img_feature = extract_img_feature(pil_img)

        img_byte_arr = io.BytesIO()
        pil_img.save(img_byte_arr, format='PNG')
        img_byte_arr = img_byte_arr.getvalue()
        img_base64_str = base64.b64encode(img_byte_arr).decode()
        tem = [uuid.uuid4().hex, row[0], img_base64_str, row[2], row[3], row[4], row[5], img_feature]
        item_list.append(tem)
        print(row_idx)
    insert_info_to_db(item_list)


def insert_info_to_db(item_list):
    conn = sqlite3.connect('data/note.db')
    # 我们也可以使用上下文管理器，这样可以在完成后自动提交和关闭连接：
    with conn:
        cur = conn.cursor()
        cur.executemany('INSERT INTO note_word VALUES (?,?,?,?,?,?,?,?)', item_list)


if __name__ == '__main__':
    load_excel_file()
